import torch
import matplotlib.pyplot as plt
import numpy as np

# 读取数据
data = torch.load("dataset_1dot25_continus_walk.pt", weights_only=False)

print("数据类型:", type(data))
print("样本数量:", len(data))

# 取一个样本（可修改索引）
sample = data[10]
src = sample["src"] if isinstance(sample, dict) else sample[0]
tgt = sample["tgt"] if isinstance(sample, dict) else sample[1]

src = np.array(src)
tgt = np.array(tgt)

print("src shape:", src.shape)
print("tgt shape:", tgt.shape)

# 不同的时间轴
time_src = np.arange(src.shape[0])
time_tgt = np.arange(tgt.shape[0])

plt.figure(figsize=(12, 14))

# 绘制 src
num_features = src.shape[1]
for i in range(num_features):
    ax1 = plt.subplot(5, 2, i + 1)
    ax1.plot(time_src, src[:, i], 'b-o', label=f"src feature {i}", linewidth=1.5, markersize=4)
    ax1.set_xlabel("Timestep")
    ax1.set_ylabel("Src Value", color='b')
    ax1.tick_params(axis='y', labelcolor='b')
    ax1.grid(True, linestyle='--', alpha=0.6)
    ax1.set_title(f"Feature {i}")
    ax1.legend()

# 绘制 tgt（第10个子图）
ax2 = plt.subplot(5, 2, 10)
ax2.plot(time_tgt, tgt[:, 0], 'r-o', label="Target", linewidth=1.5, markersize=4)
ax2.set_xlabel("Timestep")
ax2.set_ylabel("Target Value", color='r')
ax2.tick_params(axis='y', labelcolor='r')
ax2.grid(True, linestyle='--', alpha=0.6)
ax2.set_title("Target")
ax2.legend()

plt.tight_layout()
plt.show()
